package storage

import (
	"context"
	"encoding/json"
	"log/slog"
	"strings"
	"time"

	"log"

	"github.com/dgraph-io/dgo/v240"
	"github.com/dgraph-io/dgo/v240/protos/api"
	"google.golang.org/grpc"
	"google.golang.org/grpc/credentials/insecure"
)

type CancelFunc func()

type BaseRelation struct {
	Uid         string   `json:"uid,omitempty"`
	Name        string   `json:"name,omitempty"`
	RelDesc     string   `json:"rel_desc,omitempty"`
	RelKeyWords string   `json:"rel_keyword,omitempty"`
	SrcName     string   `json:"src_name,omitempty"`
	TgtName     string   `json:"tgt_name,omitempty"`
	Weight      float64  `json:"weight,omitempty"`
	DType       []string `json:"dgraph.type,omitempty"`
	DocId       string   `json:"doc_id,omitempty"`
	Rank        int      `json:"rank,omitempty"`
}

type Relation struct {
	BaseRelation
	Owes []Entity `json:"owe,omitempty"`
}

type BaseEntity struct {
	Uid            string   `json:"uid,omitempty"`
	Name           string   `json:"name,omitempty"`
	DType          []string `json:"dgraph.type,omitempty"`
	Kinds          []string `json:"kind,omitempty"`
	EntityDesc     string   `json:"entity_desc,omitempty"`
	EntityKeyWords string   `json:"entity_keyword,omitempty"`
	// Owns           []Relation `json:"owns,omitempty"`
	DocId   string `json:"doc_id,omitempty"`
	ChunkId string `json:"chunk_id,omitempty"`
	Rank    int    `json:"rank,omitempty"`
}

type Entity struct {
	BaseEntity
	// Associations []Entity `json:"associations,omitempty"`
	Owns []Relation `json:"own,omitempty"`
}

type ReverseEntity struct {
	BaseEntity
	Owes []Relation `json:"~owe,omitempty"`
}

type RelationEntity struct {
	BaseRelation
	Owns []BaseEntity `json:"~own,omitempty"`
	Owes []BaseEntity `json:"owe,omitempty"`
}

type DgraphService struct {
	dg     *dgo.Dgraph
	cancel CancelFunc
	ctx    context.Context
}

type GraphDBConfig struct {
	DBType   string `yaml:"db_type"`
	Url      string `yaml:"url"`
	Username string `yaml:"username"`
	Password string `yaml:"password"`
	NeedAuth bool   `yaml:"need_auth"`
}

func NewDgraphService(ctx context.Context, config *GraphDBConfig) (*DgraphService, error) {
	conn, err := grpc.NewClient(config.Url, grpc.WithTransportCredentials(insecure.NewCredentials()))
	if err != nil {
		log.Fatal("While trying to dial gRPC")
	}

	dc := api.NewDgraphClient(conn)
	dg := dgo.NewDgraphClient(dc)

	// Perform login call. If the Dgraph cluster does not have ACL and
	// enterprise features enabled, this call should be skipped.
	if config.NeedAuth {
		for {
			// Keep retrying until we succeed or receive a non-retriable error.
			err = dg.Login(ctx, config.Username, config.Password)
			if err == nil || !strings.Contains(err.Error(), "Please retry") {
				break
			}
			time.Sleep(time.Second)
		}
		if err != nil {
			slog.Error("While trying to login", slog.Any("error", err.Error()))
		}

	}

	service := &DgraphService{dg: dg, ctx: ctx}
	service.cancel = func() {
		if err := conn.Close(); err != nil {
			log.Printf("Error while closing connection:%v", err)
		}
	}

	return service, nil
}

func (s *DgraphService) SetDocSchema() error {
	op := &api.Operation{}
	// op.Schema = `
	// 	name: string @index(exact) .
	// 	src_name: string @index(exact) .
	// 	tgt_name: string @index(exact) .
	// 	entity_desc: string .
	// 	entity_keyword: string .
	// 	own: [uid] @reverse .
	// 	rel_desc: string .
	// 	rel_keyword: string .
	// 	doc_id: string @index(exact) .
	// 	chunk_id: string @index(exact) .
	// 	owe: uid .
	// `
	op.Schema = `
		name: string @index(exact) .
		src_name: string @index(exact) .
		tgt_name: string @index(exact) .
		entity_desc: string .
		entity_keyword: string .
		own: [uid]@reverse  .
		rel_desc: string .
		kind: [string] .
		rel_keyword: string .
		doc_id: string @index(exact).
		chunk_id: string @index(exact).
		owe: [uid]@reverse  .

		type Entity {
			name: string
			entity_desc: string
			kind: [string] 
			entity_keyword:string
			own: [Relation]
			doc_id: string
			chunk_id: string
		}

		type Relation {
			name: string
			rel_desc: string
		    src_name: string
			tgt_name: string
			rel_keyword: string
			doc_id: string
			owe: [Entity]
		}
	`
	return s.dg.Alter(s.ctx, op)
}

func (s *DgraphService) GetNodeByName(nodeName string, docId string) (*[]Entity, error) {
	q := `query Entity($nodeName: string, $docId: string) {
		entity(func: eq(name, $nodeName)) @filter(eq(doc_id,$docId)){
		    uid
			name
			dgraph.type
			entity_desc
			entity_keyword
			doc_id
			chunk_id
			own {
				uid
				name
				src_name
				tgt_name
				rel_desc
				rel_keyword
				doc_id
				owe {
				    uid
					name
					dgraph.type
					entity_desc
					entity_keyword
					doc_id
					chunk_id
				}
			}
		}
	}`
	variables := map[string]string{"$nodeName": nodeName, "$docId": docId}
	resp, err := s.dg.NewTxn().QueryWithVars(s.ctx, q, variables)
	if err != nil {
		return nil, err
	}

	type Root struct {
		Entity []Entity `json:"entity"`
	}

	var r Root
	err = json.Unmarshal(resp.Json, &r)
	if err != nil {
		return nil, err
	}
	return &r.Entity, nil
}

func (s *DgraphService) GetRelation(srcName string, tgtName string) (*[]Relation, error) {
	q := `query Entity($srcName: string, $tgtName: string) {
		relation(func: eq(src_name, $srcName)) @filter(eq(tgt_name,$tgtName)) {
				uid
				name
				src_name
				tgt_name
				rel_desc
				rel_keyword
				doc_id
			}
	}`
	variables := map[string]string{"$srcName": srcName, "$tgtName": tgtName}
	resp, err := s.dg.NewTxn().QueryWithVars(s.ctx, q, variables)
	if err != nil {
		return nil, err
	}

	type Root struct {
		Relation []Relation `json:"relation"`
	}

	var r Root
	err = json.Unmarshal(resp.Json, &r)
	if err != nil {
		return nil, err
	}
	return &r.Relation, nil
}

func (s *DgraphService) GetOwnRelation(srcName string) (*[]Relation, error) {
	q := `query Entity($srcName: string) {
		relation(func: eq(src_name, $srcName)) {
				uid
				name
				src_name
				tgt_name
				rel_desc
				rel_keyword
				doc_id
			}
	}`
	variables := map[string]string{"$srcName": srcName}
	resp, err := s.dg.NewTxn().QueryWithVars(s.ctx, q, variables)
	if err != nil {
		return nil, err
	}

	type Root struct {
		Relation []Relation `json:"relation"`
	}

	var r Root
	err = json.Unmarshal(resp.Json, &r)
	if err != nil {
		return nil, err
	}
	return &r.Relation, nil
}

func (s *DgraphService) GetReverseNode(name string, docId string) (*[]ReverseEntity, error) {
	q := `query Relation($name: string, $docId: string) {
			relation(func: eq(name,$name)) @filter(eq(doc_id,$docId)) {
						uid
						name
						dgraph.type
						entity_desc
						entity_keyword
						doc_id
						chunk_id
					~owe  {
						uid
						name
						src_name
						tgt_name
						rel_desc
						rel_keyword
					}
			}
	}`

	variables := map[string]string{"$name": name, "$docId": docId}
	resp, err := s.dg.NewTxn().QueryWithVars(s.ctx, q, variables)
	if err != nil {
		return nil, err
	}

	type Root struct {
		Relation []ReverseEntity `json:"relation"`
	}

	var r Root
	err = json.Unmarshal(resp.Json, &r)
	if err != nil {
		return nil, err
	}
	return &r.Relation, nil
}

func (s *DgraphService) GetReverseOweRelationByUid(uid string) (*[]ReverseEntity, error) {
	q := `query Relation($uid: uid) {
			relation(func: uid($uid)) {
						uid
						name
						dgraph.type
						entity_desc
						entity_keyword
						doc_id
						chunk_id
					~owe  {
						uid
						name
						src_name
						tgt_name
						rel_desc
						rel_keyword
					}
			}
	}`

	variables := map[string]string{"$uid": uid}
	resp, err := s.dg.NewTxn().QueryWithVars(s.ctx, q, variables)
	if err != nil {
		return nil, err
	}

	type Root struct {
		Relation []ReverseEntity `json:"relation"`
	}

	var r Root
	err = json.Unmarshal(resp.Json, &r)
	if err != nil {
		return nil, err
	}
	return &r.Relation, nil
}

func (s *DgraphService) GetOweRelation(tgtName string) (*[]Relation, error) {
	q := `query Entity($tgtName: string) {
		relation(func: eq(tgt_name, $tgtName)) {
				uid
				src_name
				tgt_name
				rel_desc
				rel_keyword
				owe  {
				  uid
				  name
				  dgraph.type
				  entity_desc
				  entity_keyword
				  doc_id
				  chunk_id
				}
			}
	}`

	variables := map[string]string{"$tgtName": tgtName}
	resp, err := s.dg.NewTxn().QueryWithVars(s.ctx, q, variables)
	if err != nil {
		return nil, err
	}

	type Root struct {
		Relation []Relation `json:"relation"`
	}

	var r Root
	err = json.Unmarshal(resp.Json, &r)
	if err != nil {
		return nil, err
	}

	return &r.Relation, nil
}

func (s *DgraphService) GetRelationEntity(srcName string, tagName string, docId string) (*[]RelationEntity, error) {
	q := `query Entity($src: string, $tgt: string, $docId: string) {
		relation(func: eq(tgt_name,$tgt))  @filter(eq(src_name,$src))  {
			uid
			name
			src_name
			tgt_name
			rel_desc
			doc_id
			rel_keyword 
			owe  @filter(eq(doc_id,$docId)){
					uid
					name
					dgraph.type
					entity_desc
					entity_keyword
					doc_id
					chunk_id
			     }
           ~own  @filter(eq(doc_id,$docId)){
					uid
					name
					dgraph.type
					entity_desc
					entity_keyword
					doc_id
					chunk_id
				}
			}
	}`
	variables := map[string]string{"$src": srcName, "$tgt": tagName, "$docId": docId}
	resp, err := s.dg.NewTxn().QueryWithVars(s.ctx, q, variables)
	if err != nil {
		return nil, err
	}

	type Root struct {
		Relation []RelationEntity `json:"relation"`
	}

	var r Root
	err = json.Unmarshal(resp.Json, &r)
	if err != nil {
		return nil, err
	}
	return &r.Relation, nil
}

func (s *DgraphService) GetRelationEntityByUid(uid string) (*[]RelationEntity, error) {
	q := `query Entity($uid: string) {
		relation(func: uid($uid)){
			uid
			name
			src_name
			tgt_name
			rel_desc
			doc_id
			rel_keyword 
			owe {
					uid
					name
					dgraph.type
					entity_desc
					entity_keyword
					doc_id
					chunk_id
			    }
           ~own {
					uid
					name
					dgraph.type
					entity_desc
					entity_keyword
					doc_id
					chunk_id
				}
			}
	}`
	variables := map[string]string{"$uid": uid}
	resp, err := s.dg.NewTxn().QueryWithVars(s.ctx, q, variables)
	if err != nil {
		return nil, err
	}

	type Root struct {
		Relation []RelationEntity `json:"relation"`
	}

	var r Root
	err = json.Unmarshal(resp.Json, &r)
	if err != nil {
		return nil, err
	}
	return &r.Relation, nil
}

type OwnEdge struct {
	Uid  string            `json:"uid"`
	Owns map[string]string `json:"own"`
}

type OweEdge struct {
	Uid  string            `json:"uid"`
	Owes map[string]string `json:"owe"`
}

func (s *DgraphService) SetOwnEdge(srcId string, relateId string) error {
	association := OwnEdge{
		Uid: srcId,
		Owns: map[string]string{
			"uid": relateId,
		},
	}

	eb, err := json.Marshal(association)
	if err != nil {
		return err
	}

	mu := &api.Mutation{
		CommitNow: true,
		SetJson:   eb,
	}

	_, err = s.dg.NewTxn().Mutate(s.ctx, mu)
	if err != nil {
		return err
	}

	return nil
}

func (s *DgraphService) SetOweEdge(tgtId string, relateId string) error {
	association := OweEdge{
		Uid: relateId,
		Owes: map[string]string{
			"uid": tgtId,
		},
	}

	eb, err := json.Marshal(association)
	if err != nil {
		return err
	}

	mu := &api.Mutation{
		CommitNow: true,
		SetJson:   eb,
	}

	_, err = s.dg.NewTxn().Mutate(s.ctx, mu)
	if err != nil {
		return err
	}

	return nil
}

func (s *DgraphService) GetNodeByUid(uid string) (*Entity, error) {
	q := `query Entity($uid: string) {
		entity(func: uid($uid)) {
		    uid
			name
			dgraph.type
			entity_desc
			entity_keyword
			doc_id
			chunk_id
			own {
				uid
				src_name
				tgt_name
				rel_desc
				rel_keyword
				owe {
				    uid
					name
					dgraph.type
					entity_desc
					entity_keyword
					doc_id
					chunk_id
				}
			}
		}
	}`
	variables := map[string]string{"$uid": uid}
	resp, err := s.dg.NewTxn().QueryWithVars(s.ctx, q, variables)
	if err != nil {
		return nil, err
	}

	type Root struct {
		Entity []Entity `json:"entity"`
	}

	var r Root
	err = json.Unmarshal(resp.Json, &r)
	if err != nil {
		return nil, err
	}
	return &r.Entity[0], nil
}

func (s *DgraphService) UpsertNode(entity *BaseEntity) error {
	eb, err := json.Marshal(entity)
	if err != nil {
		return err
	}

	mu := &api.Mutation{
		CommitNow: true,
		SetJson:   eb,
	}

	response, err := s.dg.NewTxn().Mutate(s.ctx, mu)
	if err != nil {
		return err
	}

	if len(response.Uids) > 0 {
		entity.Uid = response.Uids[entity.Name]
	}
	return nil
}

func (s *DgraphService) UpsertRelation(relation *BaseRelation) error {
	rb, err := json.Marshal(relation)
	if err != nil {
		return err
	}

	mu := &api.Mutation{
		CommitNow: true,
		SetJson:   rb,
	}

	response, err := s.dg.NewTxn().Mutate(s.ctx, mu)
	if err != nil {
		return err
	}
	for _, id := range response.Uids {
		relation.Uid = id
		break
	}
	return nil
}

func (s *DgraphService) DeleteDoc(docId string) error {
	entityList, err := s.getNodeByDocId(docId)
	if err != nil {
		return err
	}

	for _, entity := range *entityList {
		err = s.DeleteNode(entity.Uid)
		if err != nil {
			return err
		}
	}

	relationList, err := s.GetRelationByDoc(docId)
	if err != nil {
		return err
	}

	for _, relation := range *relationList {
		err = s.DeleteNode(relation.Uid)
		if err != nil {
			return err
		}
	}
	return nil
}

func (s *DgraphService) GetRelationByDoc(docId string) (*[]Relation, error) {
	q := `query Entity($srcName: string, $tgtName: string) {
		relation(func: eq(doc_id, $docId)) {
				uid
				name
				src_name
				tgt_name
				rel_desc
				rel_keyword
				doc_id
			}
	}`
	variables := map[string]string{"$docId": docId}
	resp, err := s.dg.NewTxn().QueryWithVars(s.ctx, q, variables)
	if err != nil {
		return nil, err
	}

	type Root struct {
		Relation []Relation `json:"relation"`
	}

	var r Root
	err = json.Unmarshal(resp.Json, &r)
	if err != nil {
		return nil, err
	}
	return &r.Relation, nil
}

func (s *DgraphService) getNodeByDocId(docId string) (*[]Entity, error) {
	q := `query Entity($docId: string) {
		entity(func: eq(doc_id, $docId)) {
		    uid
			name
			dgraph.type
			entity_desc
			entity_keyword
			doc_id
			chunk_id
			own {
				uid
				name
				src_name
				tgt_name
				rel_desc
				rel_keyword
				doc_id
				owe {
				    uid
					name
					dgraph.type
					entity_desc
					entity_keyword
					doc_id
					chunk_id
				}
			}
		}
	}`
	variables := map[string]string{"$docId": docId}
	resp, err := s.dg.NewTxn().QueryWithVars(s.ctx, q, variables)
	if err != nil {
		return nil, err
	}

	type Root struct {
		Entity []Entity `json:"entity"`
	}

	var r Root
	err = json.Unmarshal(resp.Json, &r)
	if err != nil {
		return nil, err
	}
	return &r.Entity, nil
}

func (s *DgraphService) DeleteNode(uid string) error {
	d := map[string]string{"uid": uid}

	pb, err := json.Marshal(d)
	if err != nil {
		return err
	}

	mu := &api.Mutation{
		CommitNow:  true,
		DeleteJson: pb,
	}

	response, err := s.dg.NewTxn().Mutate(s.ctx, mu)
	if err != nil {
		return err
	}

	slog.Debug("delete node rsp:", slog.Any("response:", response.Json))
	return nil
}

func (s *DgraphService) BatchDeleteOweEdge(relationUids []string, eid string) error {
	reqs := make([]*OweEdge, 0)
	for _, uid := range relationUids {
		req := &OweEdge{
			Uid: uid,
			Owes: map[string]string{
				"uid": eid,
			},
		}
		reqs = append(reqs, req)
	}

	pb, err := json.Marshal(reqs)
	if err != nil {
		return err
	}

	mu := &api.Mutation{
		CommitNow:  true,
		DeleteJson: pb,
	}

	response, err := s.dg.NewTxn().Mutate(s.ctx, mu)
	if err != nil {
		return err
	}
	slog.Debug("delete node rsp:", slog.Any("response:", response.Json))
	return nil
}

func (s *DgraphService) BatchDeleteOwnEdge(relationUids []string, eid string) error {
	reqs := make([]*OwnEdge, 0)
	for _, uid := range relationUids {
		req := &OwnEdge{
			Uid: uid,
			Owns: map[string]string{
				"uid": eid,
			},
		}
		reqs = append(reqs, req)
	}

	pb, err := json.Marshal(reqs)
	if err != nil {
		return err
	}

	mu := &api.Mutation{
		CommitNow:  true,
		DeleteJson: pb,
	}

	response, err := s.dg.NewTxn().Mutate(s.ctx, mu)
	if err != nil {
		return err
	}
	slog.Debug("delete node rsp:", slog.Any("response:", response.Json))
	return nil
}

func (s *DgraphService) DeleteRelation(srcName string, tgtName string) error {

	rels, err := s.GetRelation(srcName, tgtName)
	if err != nil {
		return err
	}

	if rels == nil {
		return nil
	}

	if len(*rels) < 1 {
		return nil
	}

	for _, r := range *rels {
		d := map[string]string{"uid": r.Uid}
		pb, err := json.Marshal(d)
		if err != nil {
			return err
		}
		mu := &api.Mutation{
			CommitNow:  true,
			DeleteJson: pb,
		}
		response, err := s.dg.NewTxn().Mutate(s.ctx, mu)
		if err != nil {
			return err
		}

		slog.Info("delete relation rsp:", slog.Any("response:", response.Json))
		return nil
	}
	return nil
}
